#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@IDE     ：PyCharm 
@Author  ：kiway
@Date    ：2024/7/26 14:12 
"""

import asyncio
import multiprocessing
import traceback
from asyncio import Queue as aQueue
from ..tools import split_list

from .BaseEngine import BaseEngine, BarEvent, OrderBookEvent, TickEvent, TradeEvent, AggTradeEvent, EVENT_SWITCH

from ..logger import pretty_print
import ccxt.pro

from copy import deepcopy
from ..config import config
from ..tools import get_base_symbol_name
import pandas as pd
import datetime
from .RowAsyncQ import Subscriber


class MarketEngine(BaseEngine):

    def __init__(self, config_task, future_ls, spot_ls, main_engine):
        super().__init__("行情引擎", main_engine)
        self.task_config = config_task
        self.future_ls = future_ls
        self.spot_ls = spot_ls

    def _start(self):

        spot_tasks_ls = split_list(self.spot_ls, config.max_num_pipe)
        future_tasks_ls = split_list(self.future_ls, config.max_num_pipe)

        count = 1
        for spot_ls, future_ls in zip(spot_tasks_ls, future_tasks_ls):
            p = multiprocessing.Process(target=self._run, args=(count, spot_ls, future_ls, self.task_config))
            p.start()
            count += 1

    async def on_event(self, strategy_o_dict):
        inner_msg_q = Subscriber()
        while True:
            msg = await inner_msg_q.get()

            try:
                if msg == EVENT_SWITCH:
                    for symbol, strategy in strategy_o_dict.items():
                        strategy.switch_time_cut()

            except:
                traceback.print_exc()

    def _run(self, count, spot_ls, future_ls, task_config):
        """行情模块启动主要逻辑
        批次订阅交易所的数据，使用异步获取数据
        单个队列无法处理过多消息，这里进行分批
        """

        loop = asyncio.new_event_loop()
        strategy_o_dict = {}

        for symbols in spot_ls + future_ls:
            st = task_config["strategy"](symbols)
            st.set_main_engine(self.main_engine)
            strategy_o_dict[symbols] = st

        for k, symbol_ls in enumerate([spot_ls, future_ls]):
            if k == 0:
                exchange = task_config["exchange"]
                exchange_options_defaultType = "spot"
            else:
                exchange = task_config["exchange"]
                exchange_options_defaultType = "future"

            _bar_istrue = task_config["is_callback"]["bar_istrue"]
            _orderbook_istrue = task_config["is_callback"]["orderbook_istrue"]
            _tick_istrue = task_config["is_callback"]["tick_istrue"]
            _allbar_istrue = task_config["is_callback"]["allbar_istrue"]
            _aggTrade_istrue = task_config["is_callback"]["aggTrade_istrue"]
            freq = "1m"

            batch_spot = BatchSpot(symbol=symbol_ls,
                                   exchange_config=exchange,
                                   exchange_options=exchange_options_defaultType,
                                   freq=freq,
                                   bar_istrue=_bar_istrue,
                                   orderbook_istrue=_orderbook_istrue,
                                   tick_istrue=_tick_istrue,
                                   allbar_istrue=_allbar_istrue,
                                   aggTrade_istrue=_aggTrade_istrue,
                                   main_engine=self.main_engine,
                                   strategy=strategy_o_dict
                                   )

            loop.create_task(batch_spot.bar_run())
            loop.create_task(batch_spot.tick_run())
            loop.create_task(batch_spot.orderbook_run())
            loop.create_task(batch_spot.agg_trade_run())

        loop.create_task(self.on_event(strategy_o_dict))
        # loop.create_task(self._run_timer(count))
        loop.run_forever()

    async def _run_timer(self, num):
        """定时log行情模块"""
        while True:
            await asyncio.sleep(100)
            pretty_print(f"行情模块-{num}-心跳...")


class BatchSpot:

    def __init__(self, symbol, exchange_config, exchange_options, freq, bar_istrue, orderbook_istrue,
                 tick_istrue, allbar_istrue, aggTrade_istrue, main_engine, strategy):

        self.freq = freq
        self.symbol = symbol
        self.bar_istrue = bar_istrue
        self.orderbook_istrue = orderbook_istrue
        self.tick_istrue = tick_istrue
        self.allbar_istrue = allbar_istrue
        self.aggTrade_istrue = aggTrade_istrue
        self.main_engine = main_engine
        self.config = exchange_config
        self.exchange_options = exchange_options
        self.exchange = None
        self.strategy_o_dict = strategy
        self._init_exchange()

    def _init_exchange(self):
        self.exchange = getattr(ccxt.pro, self.config['id'])({
            'options': {"defaultType": self.exchange_options}
        })

        self.exchange.set_sandbox_mode(self.config["sandbox_mode"])

    #
    async def bar_run(self):
        """
        这个就是子进程运行的函数，接收任务列表和用于进程间通讯的Queue
        :param  [{exchange:,  symbol:}]

        """

        if self.bar_istrue or self.allbar_istrue:
            await self._on_bar(self.exchange, self.symbol)

    async def tick_run(self):
        """
        这个就是子进程运行的函数，接收任务列表和用于进程间通讯的Queue
        :param  [{exchange:,  symbol:}]

        """
        if self.tick_istrue:
            await self._on_tick(self.exchange, self.symbol)

    async def orderbook_run(self):
        """
        这个就是子进程运行的函数，接收任务列表和用于进程间通讯的Queue
        :param  [{exchange:,  symbol:}]

        """
        if self.orderbook_istrue:
            await self._on_orderbook(self.exchange, self.symbol)

    async def agg_trade_run(self):
        if self.aggTrade_istrue or self.bar_istrue or self.allbar_istrue:
            await self._on_aggTrade(self.exchange, self.symbol)

    async def _on_bar(self, exchange, symbol):

        while True:
            try:
                bars = await exchange.watch_ohlcv_for_symbols([[s, self.freq] for s in symbol])
                symbol_str = list(bars.keys())[0]
                self.strategy_o_dict[symbol_str].on_bar(symbol_str, bars)
                self.main_engine.put_market_data_event(BarEvent(symbol=symbol_str, exchange=exchange.id, data=bars))

            except:
                traceback.print_exc()

    async def _on_tick(self, exchange, symbol):

        while True:
            try:
                ticks = await exchange.watch_tickers(symbol)
                symbol_str = list(ticks.keys())[0]
                self.strategy_o_dict[symbol_str].on_tick(symbol_str, ticks)

            except:
                traceback.print_exc()

    async def _on_orderbook(self, exchange, symbol):

        while True:
            try:
                orderbooks = await exchange.watch_order_book_for_symbols(symbol)
                symbol_str = orderbooks["symbol"]
                self.strategy_o_dict[symbol_str].on_orderbook(orderbooks["symbol"], orderbooks)

            except:
                traceback.print_exc()

    async def _on_aggTrade(self, exchange, symbol):
        # print("推送aggtrade")
        while True:
            try:
                agg_trade = await exchange.watch_trades_for_symbols(symbol, params={"name": "aggTrade"})

                symbol_str = agg_trade[0]["symbol"]
                # print(agg_trade)
                if self.allbar_istrue or self.bar_istrue:
                    self.strategy_o_dict[symbol_str]._generate_bar_from_aggtrade(symbol_str, agg_trade)

                if self.aggTrade_istrue:
                    self.strategy_o_dict[symbol_str].on_aggTrade(symbol_str, agg_trade)

                self.main_engine.put_market_data_event(
                    AggTradeEvent(symbol=symbol_str, exchange=exchange.id, data=agg_trade))

            except:
                traceback.print_exc()
